### Mantel-Haenszel extension test
mantelhaen.ext.test <- function(tbl2k){
  d <- dim(x)
  s <- d[3] # s层
  b <- d[1] # 疾病的二态，b=2
  k <- d[2] # 暴露变量的k个水平
  score <- 1:k # 评分
  O <- sum(x[2, , 1:s]*score)
  Ni <- array(0, s) # 第s层总和
  ni <- array(0, s) # 第s层第二行边际和
  mi <- array(0, s) # 第s层第一行边际和
  for(i in 1:s){
    Ni[i] <- sum(x[ , , i])
    ni[i] <- sum(x[2, , i])
    mi[i] <- sum(x[1, , i])
  }
  
  s1 <- colSums(colSums(x[ , , 1:s])*score)
  s2 <- colSums(colSums(x[ , , 1:s])*score^2)
  E <- sum(s1*ni/Ni)
  V <- sum(ni*mi*(Ni*s2 - s1^2)/(Ni^2*(Ni-1)))
  X <- (abs(O-E)-0.5)^2 / V
  p <- 1 - pchisq(q=X, df=1)
  cat("Chi-square = ", X, ", df = 1 , p-value =", p, "\n")
  res <- list(statistics=X, df=1, p.value=p)
}

### MH- stratified analysis
mhor <- function(..., mhtable=NULL, decimal=2, graph=TRUE, design="cohort") {
  if(is.null(mhtable)) {mhtable <- table(...)}else{mhtable <- as.table(mhtable)}
  a <-0
  A <-0
  Vara <-0
  numerator <- 0
  denominator <-0
  or <- c(1:dim(mhtable)[3])
  logse <- c(1:dim(mhtable)[3])
  lowlim <- c(1:dim(mhtable)[3])
  uplim <- c(1:dim(mhtable)[3])
  p.value <- c(1:dim(mhtable)[3])
  stratlab <- levels(as.data.frame(mhtable)[,3]) # Vector labelling strata
  tabodds <- c(1:(4*length(stratlab)))
  dim(tabodds) <- c(length(stratlab), 4)
  p <-0; q <-0; r <-0; s <-0; pr <-0; ps <-0; qr <-0; qs <-0; psqr <-0
  for (i in 1:dim(mhtable)[3]) 
  {
    # OR, ln(SE) and 95 ci for each staratum
    or[i] <- fisher.test(as.table(mhtable[,,i]))$estimate
    lowlim[i] <- fisher.test(as.table(mhtable[,,i]))$conf.int[1]
    uplim[i] <- fisher.test(as.table(mhtable[,,i]))$conf.int[2]
    p.value[i] <- fisher.test(as.table(mhtable[,,i]))$p.value
    # Computing MH odds ratio and standard error
    numerator <- numerator+ mhtable[1,1,i]*mhtable[2,2,i]/sum(mhtable[,,i])
    denominator <- denominator+mhtable[1,2,i]*mhtable[2,1,i]/sum(mhtable[,,i])
    p <- p+(mhtable[1,1,i]+mhtable[2,2,i])/sum(mhtable[,,i])
    q <- q+(mhtable[1,2,i]+mhtable[2,1,i])/sum(mhtable[,,i])
    r <- numerator
    s <- denominator
    pr <- pr+(mhtable[1,1,i]+mhtable[2,2,i])/sum(mhtable[,,i])*
      mhtable[1,1,i]*mhtable[2,2,i]/sum(mhtable[,,i])
    ps <- ps+(mhtable[1,1,i]+mhtable[2,2,i])/sum(mhtable[,,i])*
      mhtable[1,2,i]*mhtable[2,1,i]/sum(mhtable[,,i])
    qr <- qr+(mhtable[1,2,i]+mhtable[2,1,i])/sum(mhtable[,,i])*
      mhtable[1,1,i]*mhtable[2,2,i]/sum(mhtable[,,i])
    qs <- qs+(mhtable[1,2,i]+mhtable[2,1,i])/sum(mhtable[,,i])*
      mhtable[1,2,i]*mhtable[2,1,i]/sum(mhtable[,,i])
    psqr <- psqr+(mhtable[1,1,i]+mhtable[2,2,i])/sum(mhtable[,,i])*
      mhtable[1,2,i]*mhtable[2,1,i]/sum(mhtable[,,i])+
      (mhtable[1,2,i]+mhtable[2,1,i])/sum(mhtable[,,i])*
      mhtable[1,1,i]*mhtable[2,2,i]/sum(mhtable[,,i])
    # Computing chi-squared
    a <- a+ mhtable[1,1,i]
    A <- A+sum(mhtable[,1,i])*sum(mhtable[1,,i])/sum(mhtable[,,i])
    Vara <- Vara +  sum(mhtable[, 1, i]) / (sum(mhtable[, , i]) - 1) * 
      sum(mhtable[1, , i]) * 
      sum(mhtable[, 2, i]) * 
      sum(mhtable[2, , i]) /
      sum(mhtable[, , i])^2
    
    # Individual stratum
    tabodds[i,] <- c(or[i], lowlim[i], uplim[i], p.value[i])
  }
  cat("\n")
  colnames(tabodds) <- c("OR", "lower lim.", "upper lim.", "P value")
  collab <- colnames(as.data.frame(mhtable))[3]
  cat("Stratified analysis by ",(collab), "\n")
  rownames(tabodds) <- paste(collab, stratlab, "")
  mhor <- numerator/denominator
  mhlogse <- sqrt(pr/2/r^2 + psqr/2/r/s + qs/2/s^2)
  mhlolim <- exp(log(mhor)-qnorm(0.975)*mhlogse)
  mhhilim <- exp(log(mhor)+qnorm(0.975)*mhlogse)
  chi2 <- abs(a-A)^2/Vara # If needs corrected chisquare: chi2 <-(abs(a-A)-1/2)^2/Vara
  mh.p.value <-pchisq(chi2,1, lower.tail=FALSE)
  het <- sum((log(or)-log(mhor))^2/(1/mhtable[1,1,]+ 1/mhtable[1,2,]+ 1/ mhtable[2,1,]+ 1/
                                      mhtable[2,2,]))
  p.value.het <- pchisq(het, length(or)-1, lower.tail=FALSE)
  tabodds1 <- rbind(tabodds, c(mhor, mhlolim, mhhilim, mh.p.value))
  rownames(tabodds1)[dim(tabodds1)[1]] <- "M-H combined"
  print(tabodds1, digit=3)
  cat("\n")
  cat("M-H Chi2(1) =", round(chi2,decimal), ", P value =", round(mh.p.value, decimal+1), "\n")
  mhresults <- list(strat.table=mhtable, mh.or=mhor, ci95=c(mhlolim, mhhilim))
  if (any(mhtable==0)){
    cat(paste("\n","One or more cells of the stratified table == 0.","\n",
              "Homogeneity test not computable.","\n","\n"))
    if(graph==TRUE){
      graph <- FALSE
      cat(paste(" Graph not drawn","\n","\n"))
    }
  }else{
    cat("Homogeneity test, chi-squared", dim(tabodds)[1]-1, "d.f. =", round(het,decimal),",",
        "P value =", round(p.value.het, decimal+1), "\n")
    cat("\n")
  }
  # mhresults
  if (graph==TRUE){
    caseexp      <- rep(0, dim(mhtable)[3])
    controlex    <- rep(0, dim(mhtable)[3])
    casenonex    <- rep(0, dim(mhtable)[3])
    controlnonex <- rep(0, dim(mhtable)[3])
    logit0       <- rep(0, dim(mhtable)[3])
    se0          <- rep(0, dim(mhtable)[3])
    logit1       <- rep(0, dim(mhtable)[3])
    se1          <- rep(0, dim(mhtable)[3])
    x            <- rep(0, 6*dim(mhtable)[3])
    y            <- rep(0, 6*dim(mhtable)[3])
    for(i in 1:dim(mhtable)[3]){
      caseexp[i] 	<- mhtable[2,2,i]
      controlex[i] 	<- mhtable[1,2,i]
      casenonex[i]	<- mhtable[2,1,i]
      controlnonex[i]	<- mhtable[1,1,i]
    }
    if(design=="case control"||design=="case-control"||design=="casecontrol"){
      for(i in 1:dim(mhtable)[3]){
        logit0[i] <- log(controlex[i]/controlnonex[i]) 
        se0[i]    <- sqrt(1/controlex[i]+1/controlnonex[i])
        logit1[i] <- log(caseexp[i]/casenonex[i])
        se1[i]    <- sqrt(1/caseexp[i]+1/casenonex[i])
        x[(1:6)+(i-1)*6] <- c(c(-1,0,1)*1.96*se0[i] + logit0[i],
                              c(-1,0,1)*1.96*se1[i] + logit1[i] )
        y[(1:6)+(i-1)*6] <- c(rep(0+0.025*(i-1),3),rep(1+0.025*(i-1),3))
      }
      
      plot(x,y, xlab="Odds of exposure",yaxt="n", xaxt="n", 
           main="Stratified case control analysis", 
           ylab=paste("Outcome=",colnames(as.data.frame(mhtable))[1],
                      ", Exposure=",colnames(as.data.frame(mhtable))[2]),pch=" ")
      for(i in 1:dim(mhtable)[3]){
        lines(x[c(1,3)+(i-1)*6],y[c(1,3)+(i-1)*6], col=i+1)
        lines(x[c(4,6)+(i-1)*6],y[c(4,6)+(i-1)*6], col=i+1)
        lines(x[c(2,5)+(i-1)*6],y[c(2,5)+(i-1)*6], col=i+1, lty=2)
        points(x[c(1,3)+(i-1)*6],y[c(1,3)+(i-1)*6], col=i+1, pch="I")
        points(x[c(4,6)+(i-1)*6],y[c(4,6)+(i-1)*6], col=i+1, pch="I")
        points(x[c(2,5)+(i-1)*6],y[c(2,5)+(i-1)*6], col=i+1,pch=22,
               cex=c(controlex[i]+controlnonex[i],caseexp[i]+casenonex[i])/sum(mhtable[,,])*8)
        text(x=(max(x)+min(x))/2, y=0.3+0.1*i, col=dim(mhtable)[3]+2-i, 
             labels=paste(collab, stratlab[dim(mhtable)[3]+1-i],": OR= ",
                          round(or[dim(mhtable)[3]+1-i],decimal)," (",round(lowlim[dim(mhtable)[3]+1-i],decimal),", ",
                          round(uplim[dim(mhtable)[3]+1-i],decimal),")",sep=""))
      }
      x1 <- exp(x)
      a <- 2^(-10:10)
      if(length(a[a>min(x1) & a<max(x1)]) >2 & length(a[a>min(x1) & a<max(x1)]) <10){
        a1 <- a[a>min(x1) & a<max(x1)]
        if(any(a1>=1))axis(1,at=log(a1[a1>=1]),labels=as.character(a1[a1>=1]))
        if(any(a1<1))axis(1,at=log(a1[a1<1]),labels=paste(as.character(1),"/",
                                                          as.character(trunc(1/a1[a1<1])), sep=""))
      }
      else
      {
        options(digit=2)
        at.x <-  seq(from=min(x),to=max(x),by=((max(x)-min(x))/5))
        labels.oddsx <- exp(at.x)
        axis(1,at=at.x, labels=as.character(round(labels.oddsx,digits=decimal)))
      }
      text(x=(max(x)+min(x))/2, y=.3, labels=paste("MH-OR"," = ",
                                                   round(mhor,decimal)," (",round(mhlolim,decimal),", ",
                                                   round(mhhilim,decimal),")",sep=""))
      text(x=(max(x)+min(x))/2, y=.2, labels=paste("homogeneity test P value"," = ",
                                                   round(p.value.het, decimal+1),sep=""))
      axis(2, at=0.025*(dim(mhtable)[3]-1)/2, labels="Control", las=1)
      axis(2, at=1+0.025*(dim(mhtable)[3]-1)/2, labels="Case", las=1)
    }
    if(design=="cohort" || design=="prospective"){
      for(i in 1:dim(mhtable)[3]){
        logit0[i] <- log(casenonex[i]/controlnonex[i]); se0[i] <-sqrt(1/casenonex[i]+1/controlnonex[i])
        logit1[i] <- log(caseexp[i]/controlex[i]); se1[i] <- sqrt(1/caseexp[i]+1/controlnonex[i])
        y[(1:6)+(i-1)*6] <- c(c(-1,0,1)*1.96*se0[i] + logit0[i],c(-1,0,1)*1.96*se1[i] + logit1[i] )
        x[(1:6)+(i-1)*6] <- c(rep(0+0.025*(i-1),3),rep(1+0.025*(i-1),3))
      }
      plot(x,y, ylab="Odds of outcome",yaxt="n", xaxt="n",
           main="Stratified prospective/X-sectional analysis",
           xlab=paste("Outcome=",colnames(as.data.frame(mhtable))[1],
                      ", Exposure=",colnames(as.data.frame(mhtable))[2]),pch=" ")
      for(i in 1:dim(mhtable)[3]){
        lines(x[(1:3)+(i-1)*6],y[(1:3)+(i-1)*6], col=i+1)
        lines(x[(4:6)+(i-1)*6],y[(4:6)+(i-1)*6], col=i+1)
        lines(x[c(2,5)+(i-1)*6],y[c(2,5)+(i-1)*6], col=i+1, lty=2)
        lines(x=c(-.02,.02)+0.025*(i-1),y=c(y[1+(i-1)*6],y[1+(i-1)*6]), col=i+1)
        lines(x=c(-.02,.02)+0.025*(i-1),y=c(y[3+(i-1)*6],y[3+(i-1)*6]), col=i+1)
        lines(x=c(.98,1.02)+0.025*(i-1),y=c(y[4+(i-1)*6],y[4+(i-1)*6]), col=i+1)
        lines(x=c(.98,1.02)+0.025*(i-1),y=c(y[6+(i-1)*6],y[6+(i-1)*6]), col=i+1)
        points(x[c(2,5)+(i-1)*6],y[c(2,5)+(i-1)*6], col=i+1, pch=22,
               cex=c((controlnonex[i]+casenonex[i]),(caseexp[i]+controlex[i]))/sum(mhtable[,,])*8)
        text(x=.5, y=0.3*(max(y)-min(y))+min(y)+ 0.1*i*(max(y)-min(y)), col=dim(mhtable)[3]+2-i, 
             labels=paste(collab,stratlab[dim(mhtable)[3]+1-i],": OR = ",
                          round(or[dim(mhtable)[3]+1-i],decimal)," (",round(lowlim[dim(mhtable)[3]+1-i],decimal),", ",
                          round(uplim[dim(mhtable)[3]+1-i],decimal),")",sep=""))
      }
      text(x=.5, y=.3*(max(y)-min(y))+min(y), labels=paste("MH-OR"," = ",
                                                           round(mhor,decimal)," (",round(mhlolim,decimal),", ",
                                                           round(mhhilim,decimal),")",sep=""))
      text(x=.5, y=.2*(max(y)-min(y))+min(y), labels=paste("homogeneity test P value"," = ",
                                                           round(p.value.het, decimal+1),sep=""))
      
      axis(1, at=0.025*(dim(mhtable)[3]-1)/2, labels="Non-exposed")
      axis(1, at=1+0.025*(dim(mhtable)[3]-1)/2, labels="Exposed")
      y1 <- exp(y)
      a <- 2^(-10:10)
      if(length(a[a>min(y1) & a<max(y1)]) >2 & length(a[a>min(y1) & a<max(y1)]) <10){
        a1 <- a[a>min(y1) & a<max(y1)]
        if(any(a1>=1)) {axis(2,at=log(a1[a1>=1]),labels=as.character(a1[a1>=1]),las=1)}
        if(any(a<1)) {axis(2,at=log(a1[a1<1]),labels=paste(as.character(1),"/",
                                                           as.character(trunc(1/a1[a1<1])), sep=""),las=1)}
      }
      else
      {
        options(digit=2)
        at.y <-  seq(from=min(y),to=max(y),by=((max(y)-min(y))/5))
        labels.oddsy <- exp(at.y)
        axis(2,at=at.y, labels=as.character(round(labels.oddsy,digits=decimal+1)),las=1)
      }
      
    }
  }
}
